import numpy as np

from rlkit.envs.panda_multitask_base import MultiTaskPickAndPlaceEnv
from . import register_env

# Copy task structure from https://github.com/jonasrothfuss/ProMP/blob/master/meta_policy_search/envs/mujoco_envs/ant_rand_goal.py
@register_env('panda-pickandplace-mass')
class PickAndPlaceMassEnv(MultiTaskPickAndPlaceEnv):
    def __init__(self, task={}, n_tasks=5, randomize_tasks=True, **kwargs):
        super(PickAndPlaceMassEnv, self).__init__(task, n_tasks, **kwargs)

    def step(self, action):
        _observation, reward, terminated, truncated, info = super().step(action)
        observation = []
        for key in _observation:
            observation.append(_observation[key])
        observation = np.concatenate(observation,axis=-1)
        return observation, reward, terminated, info

    def sample_tasks(self, num_tasks):
        # a = np.random.random(num_tasks) * 2 * np.pi
        # r = 3 * np.random.random(num_tasks) ** 0.5
        # goals = np.stack((r * np.cos(a), r * np.sin(a)), axis=-1)
        masses = [0.5,1.0,1.5,2,1.0]
        tasks = []
        for mass in masses:
            tasks.append({'goal':{'mass':mass}})
        return tasks
